package com.knxhd.listeners;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import org.springframework.stereotype.Component;

import com.alibaba.fastjson.JSONObject;
import com.knxhd.enums.Role;
import com.knxhd.pojo.LlmRequest;
import com.knxhd.pojo.LlmResponse;
import com.knxhd.pojo.LlmResponse.Choice;
import com.knxhd.pojo.LlmResponse.MessageDetail;
import com.knxhd.pojo.LlmResponse.Usage;

import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.util.RandomUtil;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.Content;
import dev.langchain4j.data.message.ContentType;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.listener.ChatModelErrorContext;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.listener.ChatModelRequestContext;
import dev.langchain4j.model.chat.listener.ChatModelResponseContext;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.ChatResponseMetadata;
import dev.langchain4j.model.openai.OpenAiChatRequestParameters;
import dev.langchain4j.model.openai.OpenAiChatResponseMetadata;
import dev.langchain4j.model.openai.OpenAiTokenUsage;
import dev.langchain4j.model.output.TokenUsage;
import lombok.extern.slf4j.Slf4j;

/**
 * Author: tianluhua tianlh@qdcares.cn
 * Date: 2025-07-13 09:12:47
 * Description: 模型监听器
 */
@Component
@Slf4j
public class CustomModelListener implements ChatModelListener {

  @Override
  public void onError(ChatModelErrorContext errorContext) {
    log.error("大模型调用异常");
    Throwable error = errorContext.error();
    error.printStackTrace();

    ChatRequest chatRequest = errorContext.chatRequest();
    System.out.println(chatRequest);

    System.out.println(errorContext.modelProvider());

    Map<Object, Object> attributes = errorContext.attributes();
    System.out.println(attributes.get("my-attribute"));
  }

  @Override
  public void onRequest(ChatModelRequestContext requestContext) {
    String sessionId = RandomUtil.randomString(64);
    requestContext.attributes().put("sessionId", sessionId);
    ChatRequest chatRequest = requestContext.chatRequest();

    List<ChatMessage> messages = chatRequest.messages();
    List<MessageDetail> messageDetails = new ArrayList<>();

    LlmRequest llmRequest = new LlmRequest();
    for (ChatMessage chatMessage : messages) {
      UserMessage userMessage = (UserMessage) chatMessage;
      List<Content> contents = userMessage.contents();
      for (Content content : contents) {
        if (ContentType.TEXT.equals(content.type())) {
          TextContent textContent = (TextContent) content;
          messageDetails.add(new MessageDetail(Role.user, textContent.text()));
        }
      }
    }

    ChatRequestParameters parameters = chatRequest
        .parameters();
    llmRequest.setMessages(messageDetails);
    llmRequest.setModelName(parameters.modelName());
    llmRequest.setTemperature(parameters.temperature());
    llmRequest.setTopP(parameters.topP());
    llmRequest.setTopK(parameters.topK());
    llmRequest.setFrequencyPenalty(parameters.frequencyPenalty());
    llmRequest.setPresencePenalty(parameters.presencePenalty());
    llmRequest.setMaxOutputTokens(parameters.maxOutputTokens());
    llmRequest.setStopSequences(parameters.stopSequences());

    System.out.println(parameters.modelName());
    System.out.println(parameters.temperature());
    System.out.println(parameters.topP());
    System.out.println(parameters.topK());
    System.out.println(parameters.frequencyPenalty());
    System.out.println(parameters.presencePenalty());
    System.out.println(parameters.maxOutputTokens());
    System.out.println(parameters.stopSequences());
    System.out.println(parameters.toolSpecifications());
    System.out.println(parameters.toolChoice());
    System.out.println(parameters.responseFormat());

    if (parameters instanceof OpenAiChatRequestParameters openAiParameters) {
      System.out.println(openAiParameters.maxCompletionTokens());
      System.out.println(openAiParameters.logitBias());
      System.out.println(openAiParameters.parallelToolCalls());
      System.out.println(openAiParameters.seed());
      System.out.println(openAiParameters.user());
      System.out.println(openAiParameters.store());
      System.out.println(openAiParameters.metadata());
      System.out.println(openAiParameters.serviceTier());
      System.out.println(openAiParameters.reasoningEffort());
    }

    System.out.println(requestContext.modelProvider());

    Map<Object, Object> attributes = requestContext.attributes();
    attributes.put("my-attribute", "my-value");

    log.info("大模型调用参数: sessionId: {}, 请求参数:{}", sessionId, JSONObject.toJSONString(llmRequest));
  }

  @Override
  public void onResponse(ChatModelResponseContext responseContext) {
    ChatResponse chatResponse = responseContext.chatResponse();

    AiMessage aiMessage = chatResponse.aiMessage();
    System.out.println(aiMessage);

    ChatResponseMetadata metadata = chatResponse.metadata();
    System.out.println(metadata.id());
    System.out.println(metadata.modelName());
    System.out.println(metadata.finishReason());

    if (metadata instanceof OpenAiChatResponseMetadata openAiMetadata) {
      System.out.println(openAiMetadata.created());
      System.out.println(openAiMetadata.serviceTier());
      System.out.println(openAiMetadata.systemFingerprint());
    }

    TokenUsage tokenUsage = metadata.tokenUsage();
    System.out.println(tokenUsage.inputTokenCount());
    System.out.println(tokenUsage.outputTokenCount());
    System.out.println(tokenUsage.totalTokenCount());
    if (tokenUsage instanceof OpenAiTokenUsage openAiTokenUsage) {
      if (openAiTokenUsage.inputTokensDetails() != null) {
        System.out.println(openAiTokenUsage.inputTokensDetails().cachedTokens());
      }
      if (openAiTokenUsage.outputTokensDetails() != null) {
        System.out.println(openAiTokenUsage.outputTokensDetails().reasoningTokens());
      }
    }

    ChatRequest chatRequest = responseContext.chatRequest();
    System.out.println(chatRequest);

    System.out.println(responseContext.modelProvider());

    Map<Object, Object> attributes = responseContext.attributes();
    System.out.println(attributes.get("my-attribute"));

    String sessionId = responseContext.attributes().get("sessionId") + "";
    LlmResponse llmResponse = LlmResponse.builder()
        .choices(ListUtil.of(new Choice(new MessageDetail(Role.getRole(aiMessage.type().name()), aiMessage.text()))))
        .usage(new Usage(tokenUsage.outputTokenCount(), tokenUsage.inputTokenCount(), tokenUsage.totalTokenCount()))
        .build();

    log.info("大模型返回结果: sessionId: {}, 请求参数:{}", sessionId, JSONObject.toJSONString(llmResponse));

  }

}
